#!/usr/bin/env python
#
# Simulated cluster components
#
# Author:   Mike Murphy <mamurph@cs.clemson.edu>
# Revision: 23 September 2009
#
#   Copyright 2009 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.

'''
Simulated cluster components, including Machine and Cluster classes. Machines
represent a single CPU core, not necessarily a single compute node (which
could have multiple cores), in a cluster. Clusters represent collections of
Machines that have associated schedulers. For simulation purposes, all
Cluster names should be unique across the simulation.
'''

from core_common import SC_CLAIMED, SC_FREE, SC_SHUTDOWN, generate_name

class Machine(object):
   '''
   Simulated cluster node. For simulation purposes, each Machine is assumed
   to represent a single CPU core, rather than a full computer. Thus, when
   utilizing external trace data with the simulation system, there may be
   a one-to-many relationship between each physical machine in the trace
   and simulated Machine objects.
   '''
   def __init__(self, name, state_callback=None, **kwds):
      '''
      @param name:           Unique name for this CPU core (str)
      @type name:            str
      @param state_callback: Callback function to invoke whenever the state
                             of the CPU changes. Two arguments are passed to
                             the callback: a reference to this machine and a
                             state code (e.g. SC_CLAIMED) from core_common.
      @param kwds:           Extra data to store in the machine description
      '''
      self.name = name
      self.claimed = False
      self.alive = True
      self.callback = state_callback
      self.extra_data = kwds
      self.current_job = None
      self.capabilities = {}
   def claim(self):
      '''
      Marks the machine as claimed for scheduling purposes
      '''
      if self.alive and not self.claimed:
         self.claimed = True
	 if self.callback:
	    self.callback(self, SC_CLAIMED)
	 return True
      return False
   def free(self):
      '''
      Marks the machine as free (unclaimed) for scheduling purposes
      '''
      self.claimed = False
      if self.callback:
         self.callback(self, SC_FREE)
   def shut_down(self, time, simobj):
      '''
      Shuts down the machine, killing any job that happens to be running at
      the time. This method is simulator-aware for use as an event.

      @param time: Current simulated time
      @type time: number
      @param simobj: Reference to the simulator
      '''
      self.alive = False
      if self.current_job is not None:
         self.current_job.kill(time, simobj)
      r = None
      if self.callback:
         r = self.callback(self, SC_SHUTDOWN)
      return r
   def adjust_length(self, length, cpu, net):
      '''
      Adjusts the length of a job based on the CPU/network consumption of the
      job. This base implementation does nothing (i.e. the length remains the
      same as the input length). Subclass Machine and override this method
      to implement overhead simulation (e.g. for virtual machines).

      @param length: Current length of the job
      @type length: number (non-negative)
      @param cpu: CPU bound of the job
      @type cpu: number in range 0-100
      @param net: Network bound of the job (100 - cpu bound)
      @type net: number in range 0-100

      @return: new length of the job (== length for this base implementation)
      '''
      return length
   def add_capability(self, cap_name, cap_obj):
      '''
      Adds a machine-specific capability to the simulated machine (e.g.
      virtualization support).

      @param cap_name: Name of the capability or service
      @type cap_name: str
      @param cap_obj: Instance object implementing the capability
      '''
      self.capabilities[cap_name] = cap_obj
   def get_capability(self, cap_name):
      '''
      @param cap_name: Name of the capability to request
      @type cap_name: str

      @return: instance of an object providing a named capability, or None
      if this machine does not provide the requested service
      '''
      obj = None
      if cap_name in self.capabilities:
         obj = self.capabilities[cap_name]
      return obj
   def remove_capability(self, cap_name):
      '''
      Removes a capability from this machine. If the capability name is not
      known, an Exception will be raised.

      @param cap_name: Name of the capability to remove
      @type cap_name: str
      '''
      del self.capabilities[cap_name]
   def __str__(self):
      result  = 'Machine:  %s\n' % self.name
      for item in self.extra_data:
         result += '   %s:  %s\n' % (item, self.extra_data[item])
      return result
#


class MachineFactory(object):
   '''
   Factory class for producing new Machines based upon named profiles. Each
   profile is associated with a class to instantiate (Machine or a subclass),
   along with keyword arguments to the class constructor.
   '''
   def __init__(self, default_class=Machine, **kwds):
      '''
      @param default_class: Default machine class to create when a profile
                            cannot be found
      @param kwds: Default keyword arguments to pass to the class
                   constructor when using the default class
      '''
      self.profiles = {}   # name: (class_name, {kwd_args})
      self.default_class = default_class
      self.default_kwds = kwds
   def add_profile(self, profile, class_ptr, **kwds):
      '''
      Adds a profile to the factory

      @param profile: Profile name
      @type profile: str
      @param class_ptr: Class to instantiate when the profile is requested
      @param kwds: Keyword arguments to the class constructor
      '''
      self.profiles[profile] = (class_ptr, kwds)
   def remove_profile(self, profile):
      '''
      Removes a profile from the factory. The profile must already be present
      in the factory, or an Exception will occur.

      @param profile: Name of the profile to remove (str)
      @type profile: str
      '''
      del self.profiles[profile]
   def new_machine(self, name, profile=''):
      '''
      Creates a new machine, using the optional profile to select an alternate
      class configuration. If the profile is omitted or cannot be found in the
      factory, an instance of the default class is constructed using the
      default keyword arguments.

      @param name: Name of the new machine
      @type name: str
      @param profile: Machine profile name
      @type profile: str

      @return: the newly instantiated machine
      '''
      m = None
      if profile in self.profiles:
         cls, kwds = self.profiles[profile]
         m = cls(name, **kwds)
      else:
         m = self.default_class(name, **self.default_kwds)
      return m
#


# A cluster contains systems and one or more schedulers (whose targets may
# overlap!). Individual machines SHOULD NOT be shared between clusters!
class Cluster(object):
   '''
   Clusters represent collections of computing resources (Machines) and one
   or more schedulers (whose targets may overlap!). Individual resources
   (Machines) generally should not be shared between clusters.
   '''
   def __init__(self, name, simobj, factory=MachineFactory(), **kwds):
      '''
      @param name: Name of this cluster. For grid-level simulations, cluster
                   names should be unique.
      @type name: str
      @param simobj: Reference to the simulator object
      @param factory: MachineFactory instance used when creating new machines
      @param kwds: Cluster extra data, which will be passed to each machine
                   created whenever the cluster is resized
      '''
      self.schedulers = {}
      self.machines = []
      self.name = name
      self.sim = simobj
      self.machine_factory = factory
      self.extra_data = kwds
      self.name_index = 1
      self.name_size = 100000
   def add_machine(self, machine):
      '''
      Adds an *existing* machine to the cluster. If a new machine needs to be
      created, use grow() instead. Note that this method does NOT add the
      machine as a target to any scheduler.

      @param machine: Machine object to add to the cluster
      '''
      if machine in self.machines:
         raise Exception('Machines must be unique')
      else:
         self.machines.append(machine)
   def add_machines(self, machines):
      '''
      Adds a list of *existing* machines to the cluster. If new machines need
      to be created, use grow() instead. Note that this method does NOT add
      any of the machines as a target to any scheduler.

      @param machines: List of Machine objects to add to the cluster
      '''
      for machine in machines:
         self.add_machine(machine)
   def remove_machine(self, machine):
      '''
      Removes a Machine object from the cluster, *without* removing the machine
      from the scheduler pool. For most purposes, use shrink() instead.

      @param machine: Machine object to remove from the cluster
      '''
      self.machines.remove(machine)
   def remove_machines(self, machines):
      '''
      Removes a list of Machine objects from the clause, *without* removing
      them from the scheduler pool. For most purposes, use shrink() instead.

      @param machines: List of Machine objects to remove from the cluster
      '''
      for machine in machines:
         self.machines.remove(machine)
   def grow(self, count, profile=''):
      '''
      Adds new machines to the cluster, provided count is greater than zero.
      Machines will be created using the machine factory specified in the
      class constructor. Note that this method does NOT add machines to any
      scheduler pool.

      @param count: Number of machines to add (positive) or subtract (negative)
                    to or from the cluster
      @type count: int
      @param profile: Optional machine profile name, for selecting machine
                      configurations
      @type profile: str

      @return: list of new Machine objects that have been added
      '''
      r = []
      basename = 'cpu.' + self.name + '-'
      for x in xrange(0, count):
         name = generate_name(basename, self.name_index, self.name_size)
         self.name_index += 1
         m = self.machine_factory.new_machine(name, profile)
         self.add_machine(m)
         r.append(m)
      return r
   def shrink(self, count, kill_busy=True):
      '''
      Shrinks the cluster by a specified number of machines. Idle machines
      will be selected for removal first, following by randomly chosen (using
      the simulator random source) busy machines (if kill_busy is True). Upon
      removal, machines will be terminated using the shut_down() method of
      the Machine object. If a request would attempt to remove more machines
      than are present in the cluster, no machines will be removed, and an
      Exception will occur.

      @param count: Number of machines to be removed
      @type count: int
      @param kill_busy: Flag to indicate whether busy machines should be
                        removed in order to reach the desired shrink count
      @type kill_busy: bool

      @return: list of Machine objects that have been removed. Note that
      fewer machines than requested may have been removed, if kill_busy is
      False.
      '''
      r = []
      if count > len(self.machines):
         raise Exception('Requested kill count exceeds cluster size')
      killed = 0
      now = self.sim.get_time()
      for m in self.machines[:]:
         if (not m.claimed) and (killed < count):
            self.remove_machine(m)
            m.shut_down(now, self.sim)
            r.append(m)
            killed += 1
      #
      if kill_busy and (killed < count):
         need = count - killed
         victims = self.sim.random.sample(self.machines.keys(), need)
         for v in victims:
            m = self.machines[v]
            self.remove_machine(m)
            m.shut_down()
            r.append(m)
            killed += 1
      #
      return r
   def add_scheduler(self, name, sched):
      '''
      Add a scheduler to the list of those known to the cluster. Note that
      this method does NOT add any machines to the scheduler pool.

      @param name: Name of the scheduler
      @type name: str
      @param sched: Scheduler object
      '''
      if name in self.schedulers:
         raise Exception('Scheduler name entries must be unique')
      else:
         self.schedulers[name] = sched
   def remove_scheduler(self, name):
      '''
      Removes a scheduler from the list of schedulers known to the cluster.
      This method does NOT remove any machines from the scheduler pool.

      @param name: Name of the scheduler
      @type name: str
      '''
      del self.schedulers[name]
   def get_scheduler(self, name=''):
      '''
      Returns a reference to a scheduler based on the scheduler name. If the
      requested name is an empty string (default), then a reference to an
      arbitrary scheduler known to the cluster will be returned. Otherwise, a
      reference to the requested scheduler will be returned, if the scheduler
      is known to the cluster. If an unknown scheduler is requested, None is
      returned.

      @param name: Name of the scheduler
      @type name: str
      '''
      sched = None
      if name == '':
         for key in self.schedulers:
            sched = self.schedulers[key]
            break
      elif name in self.schedulers:
         sched = self.schedulers[name]
      return sched
   def get_cpu_count(self):
      '''
      @rtype: int
      @return: total number of machines in the cluster
      '''
      return len(self.machines)
   def get_actual_idle_count(self):
      '''
      @rtype: int
      @return: total number of idle machines in the cluster
      '''
      count = 0
      for m in self.machines:
         if not machine.claimed:
	    for s in self.schedulers:
	       if s.is_target(m):
	          count += 1
      return count
   def get_waiting_count(self):
      '''
      @rtype: int
      @return: total number of jobs waiting on all schedulers known
               to the cluster
      '''
      count = 0
      for s in self.schedulers:
         count += s.get_waiting_count()
      return count
   def get_running_count(self):
      '''
      @rtype: int
      @return: total number of jobs running on all schedulers known
               to the cluster
      '''
      count = 0
      for s in self.schedulers:
         count += s.get_running_count()
      return count
   def get_total_queue_size(self):
      '''
      @rtype: int
      @return: total queue size (running+waiting) on all schedulers known
               to the cluster
      '''
      count = 0
      for s in self.schedulers:
         count += s.get_total_queue_size()
      return count
   def get_completed_count(self):
      '''
      @rtype: int
      @return: total number of jobs completed on all schedulers known to
      the cluster
      '''
      count = 0
      for s in self.schedulers:
         count += self.get_completed_count()
      return count
#

if __name__ == '__main__':
   from kernel import Simulator
   sim = Simulator()
   c = Cluster('test', sim)
   c.grow(10)
   c.shrink(5)
   c.grow(20)
   c.shrink(25)
#
